import os
import random
import shutil

import cv2
import numpy as np

from utils.file_util import pickleSaveData


class DataAugmentation:
    def __init__(self, data):
        self.feature = data[0]
        self.label = data[1]
        self.augNum = 10
        self.prefix = 'dataset/train'
        if os.path.exists(self.prefix):
            shutil.rmtree(self.prefix)
        os.makedirs(self.prefix)

    def augAndSaveData(self):
        count = 0
        printNum = 5000
        hintNum = 0
        dNum = self.feature.shape[0]
        for t in range(dNum):
            sinFeature = self.feature[t]
            sinLabel = self.label[t]
            f, l = self.randomCrop(sinFeature, sinLabel)
            count = self.saveSingleDataToFile(f,l,count)

            f, l = self.randomNoise(sinFeature, sinLabel)
            count = self.saveSingleDataToFile(f,l,count)

            f, l = self.randomZero(sinFeature, sinLabel)
            count = self.saveSingleDataToFile(f, l, count)

            f, l = self.randomResize(sinFeature, sinLabel)
            count = self.saveSingleDataToFile(f, l, count)

            if count // printNum > hintNum:
                print(f'current process data is {count}')
                hintNum+=1

        count = self.saveSingleDataToFile(self.feature, self.label, count)

        print(f'total train number is {count}')

    def saveSingleDataToFile(self, feature, label, count):
        feature = np.array(feature)
        label = np.array(label)

        sampleLen = label.shape[0]
        for i in range(sampleLen):
            savePath = os.path.join(self.prefix, f'{count}.dp')
            data = [feature[i], label[i]]
            pickleSaveData(data, savePath)
            count+=1
        return count

    def augAllData(self):
        fFeature = self.feature.copy()
        fLabel = self.label.copy()
        dNum = self.feature.shape[0]
        for t in range(dNum):
            sinFeature = self.feature[t]
            sinLabel = self.label[t]
            f, l = self.randomCrop(sinFeature, sinLabel)
            fFeature = np.concatenate((fFeature, f), axis=0)
            fLabel = np.concatenate((fLabel, l), axis=0)

            # fn, ln = self.randomNoise(sinFeature, sinLabel)
            # fFeature = np.concatenate((fFeature, fn), axis=0)
            # fLabel = np.concatenate((fLabel, ln), axis=0)

        return [fFeature, fLabel]

    def randomCrop(self, sinFeature, sinLabel, offPad=40):
        # offList = random.sample(range(-offPad, offPad), self.augNum)
        # 左加右减
        offList = random.sample(range(-offPad, -1), self.augNum)
        shape = sinFeature.shape
        channel = shape[0]
        resFeature = []
        resLabel = []
        for offNum in offList:
            cData = []
            for t in range(channel):
                curImg = sinFeature[t]
                pad = np.zeros((abs(offNum), shape[2]))
                res = curImg
                if(offNum > 0):
                    res = np.concatenate((curImg[offNum:, :], pad), axis=0)
                if(offNum < 0):
                    res = np.concatenate((pad, curImg[0:offNum ,:]), axis=0)
                cData.append(res)
            resFeature.append(cData)
            resLabel.append(sinLabel)
        # resFeature = np.array(resFeature)
        # resLabel = np.array(resLabel)
        return resFeature, resLabel

    def randomZero(self, sinFeature, sinLabel, wid = 4):
        shape = sinFeature.shape
        channel = shape[0]
        horizonList = random.sample(range(0, shape[-2]-wid-1), self.augNum // 2)
        verticalList = random.sample(range(0, shape[-1]-wid-1), self.augNum // 2)
        resFeature = []
        resLabel = []
        for i in range(len(horizonList)):
            hData = []
            vData = []
            for t in range(channel):
                curImg = sinFeature[t]

                hNum = horizonList[i]
                hImg = curImg.copy()
                hImg[hNum:hNum+wid, :] = 0
                hData.append(hImg)

                vNum = verticalList[i]
                vImg = curImg.copy()
                vImg[:, vNum:vNum+wid] = 0
                vData.append(vImg)
            resFeature.append(hData)
            resFeature.append(vData)
            resLabel.append(sinLabel)
            resLabel.append(sinLabel)
        # resFeature = np.array(resFeature)
        # resLabel = np.array(resLabel)
        return resFeature, resLabel

    def randomNoise(self, sinFeature, sinLabel, div=50):
        shape = sinFeature.shape
        channel = shape[0]
        resFeature = []
        resLabel = []
        for i in range(self.augNum):
            cData = []
            noise = np.random.normal(0, 1, shape[1:]) / div
            for t in range(channel):
                curImg = sinFeature[t]
                tempImg = curImg.copy()
                tempImg += noise
                cData.append(tempImg)
            resFeature.append(cData)
            resLabel.append(sinLabel)

        # resFeature = np.array(resFeature)
        # resLabel = np.array(resLabel)
        return resFeature, resLabel

    def randomResize(self, sinFeature, sinLabel):
        shape = sinFeature.shape
        channel = shape[0]
        width = shape[1]
        height = shape[2]
        resFeature = []
        resLabel = []

        hOff = shape[1]//4
        vOff = shape[2]//4

        horizonList = random.sample(range(width-hOff, width+hOff), self.augNum)
        verticalList = random.sample(range(height-vOff, height+vOff), self.augNum)

        for i in range(len(horizonList)):
            hData = []
            vData = []
            for t in range(channel):
                curImg = sinFeature[t]

                hNum = horizonList[i]
                hImg = curImg.copy()
                hImg = cv2.resize(hImg, (height, hNum))
                hImg = self.cropOrPad(width, height, hImg)
                hData.append(hImg)

                vNum = verticalList[i]
                vImg = curImg.copy()
                vImg = cv2.resize(vImg, (vNum, width))
                vImg = self.cropOrPad(width, height, vImg)
                vData.append(vImg)

            resFeature.append(hData)
            resFeature.append(vData)
            resLabel.append(sinLabel)
            resLabel.append(sinLabel)
        # resFeature = np.array(resFeature)
        # resLabel = np.array(resLabel)
        return resFeature, resLabel

    def cropOrPad(self, width, height, img):
        shape = img.shape
        if (width <= shape[0]):
            resNum = shape[0] - width
            start = resNum // 2
            img = img[0:width, :]
        else:
            resNum = width - shape[0]
            leftNum = resNum // 2
            rightNum = resNum - leftNum
            leftPad = np.zeros( (leftNum, shape[1]) )
            rightPad = np.zeros( (rightNum, shape[1]) )
            img = np.concatenate( (leftPad, img, rightPad), axis=0)

        if (height <= shape[1]):
            resNum = shape[1] - height
            start = resNum // 2
            img = img[:, start:start+height]
        else:
            resNum = height - shape[1]
            upNum = resNum // 2
            downNum = resNum - upNum
            upPad = np.zeros( (shape[0], upNum) )
            downPad = np.zeros( (shape[0], downNum) )
            img = np.concatenate( (upPad, img, downPad), axis=1)

        return img


# img = np.ones((3, 200, 68))
# obj = DataAugmentation([[],[]])
# a, b = obj.randomNoise(img, 1)
# print(a.shape, b.shape)
